Skip to content

refactor: add backend support + cancellation and batched PBS query - #14

Merged
LDeakin merged 18 commits into
mainfrom
refactor/backend-only
Jul 30, 2026
Merged

refactor: add backend support + cancellation and batched PBS query#14
LDeakin merged 18 commits into
mainfrom
refactor/backend-only

Conversation

@LDeakin

@LDeakin LDeakin commented Jul 30, 2026

Copy link
Copy Markdown
Member

This is a stripped down version of the simple-server branch, retaining the backend support and only LocalBackend.

LLM found a few issues and I think it is much better to put the actual job tracking in the prefect worker + persistent prefect variables rather than an independent process.

Since the move to the `wex2-orchestration` persistent session,
the Prefect worker and `pbspy-server` are neighbours on the same host — the same host that
already has `qsub`/`qstat`/`qdel` on `PATH` and `/g/data` mounted locally. They talk to each
other over loopback (`127.0.0.1:9876`) using length-prefixed **pickle** frames and a shared
secret. The indirection buys nothing and costs plenty:

- **Security:** `recv_frame` calls `pickle.loads` on socket bytes *before* the API key is
  checked (`pbspy/src/pbspy/server/server.py:129`), so anyone who can reach the port has RCE
  regardless of the key. The key itself is a single global plaintext string compared with `!=`.
- **Fragility:** the server holds all job state in memory
  (`pbspy/src/pbspy/server/server.py:13`). A `supervisorctl restart pbspy-server` drops every
  wait subscription and every blocked `job.result()` raises. A `WaitRequest` for a job the
  server didn't itself submit blocks the client **forever** (no timeout).
- **No resumability:** the flow-run-local registry `_job_ids_by_run`
  (`wex2_automation/utils/gadi_pbs.py:23`) dies with the subprocess. After a restart the PBS
  job keeps running, the flow fails, the retry sees no output yet, and **resubmits a duplicate
  job**. Nothing ever reconciles against `qstat`.

Added

  • Add the public Backend interface and default LocalBackend implementation
    • JobDescription.submit() accepts a backend supplied by the caller
    • Each Job retains its backend for waiting, result retrieval, and cancellation
  • Add Job.cancel()
  • Add batched PBS state-query and deletion helpers

Changed

  • Group Job.wait_all() and Job.result_all() calls by backend so each backend can wait for its jobs together
  • Reuse a shared LocalBackend for jobs submitted without an explicit backend
  • Return None when PBS has not reported a completed job's exit status

Fixed

  • Preserve JobDescription.output_path and error_path on submitted jobs and use them when retrieving results
  • Treat cancellation of finished or unknown PBS jobs as successful

omdowley and others added 18 commits May 29, 2026 14:27
Previously pbspy could only submit jobs when running directly on a
supercomputer login node.  This adds a client/server architecture
so jobs can be submitted from any machine with SSH access.

src/pbspy/__init__.py has been split into discrete modules:

  _backend.py        - Backend Abstract Base Class
  _local_backend.py  - LocalBackend: direct qsub/qstat calls (default)
  _ssh_backend.py    - SSHBackend: talks to the daemon over SSH
  _pbs_core.py       - Core PBS operations extracted from __init__.py
  _protocol.py       - Length-prefixed pickle framing (send_frame /
                        recv_frame) and request/response dataclasses

A new top-level package pbspy_server contains:

  server.py   - TCP daemon: accepts connections, dispatches requests,
                polls qstat in a background thread, streams status
                updates to waiting clients, shuts down automatically
                when no jobs are pending and no clients are connected
  proxy.py    - SSH proxy mode: acquires a file lock, starts the daemon
                if needed, then splices stdin/stdout with the daemon
                socket so the remote client communicates transparently
  __main__.py - CLI entry point (--daemon / --proxy / --status)

pyproject.toml adds the `pbspy-server` console script.

`pbspy-server` binds a TCP port on all interfaces so that proxy
processes on other login nodes can reach it.  To prevent unauthorized
access (pickle deserialization is remote code exec), every connection
must authenticate with a shared secret token:
  - On startup the server generates secrets.token_hex(32) and writes the
    addr file on the login-node as JSON: {"host": ..., "port": ..., "token": ...}
  - Clients send AuthRequest(token=...) as the first frame; the server
    responds with AuthOkResponse or ErrorResponse and closes the connection
  - The proxy reads the token from the JSON addr file and authenticates
    using it
  - pbspy-server --status and _daemon_alive() also authenticate before
    sending PingRequest

  JobDescription.submit() now accepts an optional `backend` parameter
  (defaults to LocalBackend).  Pass an SSHBackend instance to submit to
  a remote supercomputer:
  ```python
      backend = SSHBackend("user@gadi.nci.org.au")
      job = JobDescription(...).add_command([...]).submit(backend=backend)
  ```
  Backend, LocalBackend, and SSHBackend are exported from pbspy.__all__.

Added tests to new code:
  tests/test_protocol.py
  tests/test_server.py
  tests/test_default.py
Updated pre-commit:
  .pre-commit-config.yaml: added pytest to mypy's additional_dependencies so
  the type stubs for pytest decorators are available during the hook run.
The server now runs on a machine with SSH access to the supercomputer
rather than on the login node itself.

PBS commands (qsub, qstat, file reads) are issued by the server over
SSH via the new PBSRunner. Clients connect to the server
over a plain TCP socket. A thourough cleanup has removed no-longer
required parts of the last commit.
Includes .env.example for SSH configuration. PBSPY_SSH_HOST and
PBSPY_SSH_USER are also read from environment as argparse defaults.
This fixes a bug in the Dockerfile caused by the pbspy_server module not
being installed by `hatchling`.
The Backend was being pickled when passing jobs to the server, causing
crashes because Sockets aren't pickleable. This resolves this and adds a
regression test ensuring we can pickle and unpickle a Job.
Required so that `pbs_get_result` gets output and error files correctly
if their paths are modified by the JobDescription.
Clients connecting to a server started with --api-key must send
AuthRequest as the first frame; wrong or missing key closes the
connection with ErrorResponse. Auth is opt-in: servers without
--api-key accept all connections as before.

Adds ExecRequest/ExecResponse to the protocol so clients can run
arbitrary commands through the server's SSH runner. Disabled by
default (--allow-exec / PBSPY_ALLOW_EXEC=1); a startup warning
is logged when enabled.
Introduces `pbspy-server --proxy-to HOST` (ProxyServer), a pure byte-relay
that forwards client connections to a pbspy-server running in an Gadi
persistent session via `ssh -W`. Existing `ServerBackend` clients
connect to the proxy without any changes.

Refactors shared RPC logic (auth, ping, submit/wait/result, reconnect) out of
`ServerBackend` into a new `StreamBackend` base class, leaving only the
transport (socket open/close) in the subclass.
This commit entirely removes the SSH layer from pbspy -- setting up the
transport layer is a user's concern and not this library's. The server
is simplified with the idea that it should run on a node with direct
access to PBS (e.g., an NCI Gadi persistent-session).
Retain the backend abstraction and local PBS implementation so alternate backends can be supplied by separate packages. Remove the bundled server, wire protocol, transport code, CLI, and associated documentation and tests.

Also tighten backend result typing, share the default local backend, retain legacy Mango progress-call compatibility, and document the net changes since v0.0.9.
@LDeakin
LDeakin merged commit ace54bf into main Jul 30, 2026
3 checks passed
@LDeakin
LDeakin deleted the refactor/backend-only branch July 30, 2026 00:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants